from typing import Union, Callable, Dict, List
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

import requests
import urllib3

from requests.auth import HTTPBasicAuth

import HttpHelper

urllib3.disable_warnings()


class _Range(object):
    def __init__(self, start: int, end: int, is_downloading: bool = False):
        self.start = start
        self.end = end
        self.is_downloading = is_downloading

    @staticmethod
    def create_offset_length(offset: int, length: int, is_downloading: bool = False) -> '_Range':
        return _Range(offset, offset + length - 1, is_downloading)

    @staticmethod
    def create_empty(start: int = 0):
        return _Range(start, start - 1)

    def duplicate(self):
        return _Range(self.start, self.end, self.is_downloading)

    def __len__(self):
        return self.end - self.start + 1

    def __str__(self):
        return f'[{self.start} - {self.end}]'

    # def __or__(self, other) -> Union[None, '_Range']:
    #     if self.end + 1 < other.start or other.end + 1 < self.start:
    #         return None
    #     return _Range(min(self.start, other.start), max(self.end, other.end))

    def __and__(self, other) -> Union[None, '_Range']:
        return _Range(max(self.start, other.start), min(self.end, other.end))

    # def __add__(self, other) -> '_Range':
    #     return _Range(min(self.start, other.start), max(self.end, other.end))
    #
    # def __sub__(self, other) -> '_Range':
    #     if self.empty() or other.empty():
    #         return self
    #
    #     if other.end < self.start or other.start > self.end:
    #         return self
    #
    #     if self.start < other.start:
    #         return _Range(self.start, other.start - 1)
    #
    #     if self.end > other.end:
    #         return _Range(other.end + 1, self.end)
    #
    #     return _Range(max(self.start, other.start), min(self.end, other.end))

    def empty(self) -> bool:
        return len(self) <= 0


class _DownloadRangeManager(object):
    def __init__(self,
                 offset: int,
                 length: int,
                 min_range: int = 1*1024*1024,
                 max_range: int = 10 * 1024 * 1024):
        self.lock = threading.Lock()
        self._ranges: List[_Range] = [_Range.create_offset_length(offset, length, False)]
        self.min_range = min_range
        self.max_range = max_range

    def get(self) -> _Range:
        with self.lock:
            if len(self._ranges) <= 0:
                return _Range.create_empty()

            # sort by length
            self._ranges.sort(key=lambda _r: len(_r), reverse=True)

            # try find largest idle ranges
            for i in range(len(self._ranges)):
                if not self._ranges[i].is_downloading:
                    r = self._ranges.pop(i)
                    r_len = len(r)
                    if r_len > self.max_range:
                        r_fetch = _Range.create_offset_length(r.start, self.max_range, True)
                        r_left = _Range(r_fetch.end + 1, r.end, False)
                        self._ranges.append(r_fetch)
                        self._ranges.append(r_left)
                        return r_fetch
                    else:
                        r.is_downloading = True
                        self._ranges.append(r)
                        return r

            # if no idle range, then find largest downloading range and split into two ranges
            r = self._ranges.pop(0)
            new_size = len(r) // 2
            old_size = len(r) - new_size
            if new_size < self.min_range or new_size == 0:
                self._ranges.append(r)
                return _Range.create_empty()  # don't split range if it's smaller than min_range

            # split the range
            old_range = _Range.create_offset_length(r.start, old_size, True)
            new_range = _Range.create_offset_length(old_range.end + 1, new_size, True)
            self._ranges.append(old_range)
            self._ranges.append(new_range)

            return new_range

    def commit(self, r: _Range) -> _Range:
        with self.lock:
            # find range
            original: Union[None, _Range] = None
            for i in range(len(self._ranges)):
                if self._ranges[i].start == r.start:
                    original = self._ranges.pop(i)
                    break
            if original is None:
                # something wrong, should not happen
                return _Range.create_empty(r.start)

            # generate accepted range and add left downloading range back
            accepted = r & original
            if accepted.empty():
                self._ranges.append(original)
            elif len(accepted) < len(original):
                self._ranges.append(_Range(accepted.end + 1, original.end, True))

            return accepted

    def abort(self, r: _Range) -> _Range:
        with self.lock:
            # find range
            found: Union[None, _Range] = None
            for i in range(len(self._ranges)):
                if self._ranges[i].start == r.start:
                    found = self._ranges.pop(i)
                    break
            if found is None:
                # something wrong, should not happen
                return _Range.create_empty(r.start)

            # mark range as idle
            found.is_downloading = False
            self._ranges.append(found)

            return found


class _MTContext(object):
    def __init__(self,
                 url: str,
                 auth: HTTPBasicAuth,
                 file_size: int,
                 headers: Union[None, Dict[str, str]],
                 max_thread_count: int,
                 on_receive: Callable[[int, bytes, int], None],
                 drm: _DownloadRangeManager):
        self.lock = threading.Lock()
        self.cond = threading.Condition(self.lock)
        self.url = url
        self.auth = auth
        self.downloaded_size = 0
        self.file_size = file_size
        self.headers = headers
        self.max_thread_count = max_thread_count
        self.on_receive = on_receive
        self.drm = drm
        self.downloading_tasks = 0


class Http(object):
    DEBUG: bool = False

    UA_CHROME = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_7_0)' \
                ' AppleWebKit/535.11 (KHTML, like Gecko)' \
                ' Chrome/17.0.963.56 Safari/535.11'

    HEADERS = {
        "User-Agent": UA_CHROME,
        'Accept-Encoding': 'identity'
    }

    # DOWNLOAD_BUFFER_SIZE = 512 * 1024
    DOWNLOAD_BUFFER_SIZE = 1024

    @staticmethod
    def get(url: str,
            on_receive: Union[None, Callable[[int, bytes, int], None]] = None,
            on_progress: Union[None, Callable[[int, int], None]] = None,
            min_segment_size: int = 1 * 1024 * 1024,  # min segment size, 1M as default
            max_segment_size: int = 10 * 1024 * 1024,  # max segment size, 10M as default
            max_thread_count: int = 1,
            username: str = None, password: str = None,
            headers: Union[None, Dict[str, str]] = None) -> Union[bytes, None, bool]:
        # create auth and headers
        auth = HTTPBasicAuth(username, password)
        headers = Http.get_headers(headers)

        # get file_size
        file_size = -1
        if max_thread_count > 1:
            # get file size only when multi-downloading
            response = requests.head(url, headers=headers, auth=auth)
            if 'Content-Length' in response.headers:
                content_length = response.headers['Content-Length']
                if content_length.isnumeric():
                    file_size = int(content_length)

        # receive
        if on_receive is None:
            result: bool
            buffer = []
            if file_size < 0:
                result = Http._get_single(url=url,
                                          on_receive=HttpHelper.HttpHelper.buffer_receiver(buffer),
                                          on_progress=on_progress,
                                          auth=auth,
                                          headers=headers)
            else:
                result = Http._get_mt(url=url,
                                      file_size=file_size,
                                      on_receive=HttpHelper.HttpHelper.buffer_receiver(buffer),
                                      on_progress=on_progress,
                                      min_segment_size=min_segment_size,
                                      max_thread_count=max_thread_count,
                                      auth=auth,
                                      headers=headers)

            if not result:
                return None

            buffer.sort(key=lambda item: item[0])
            bytes_list = [item[1] for item in buffer]
            return b''.join(bytes_list)

        else:
            if file_size < 0:
                return Http._get_single(url=url,
                                        on_receive=on_receive,
                                        on_progress=on_progress,
                                        auth=auth,
                                        headers=headers)
            else:
                return Http._get_mt(url=url,
                                    file_size=file_size,
                                    on_receive=on_receive,
                                    on_progress=on_progress,
                                    min_segment_size=min_segment_size,
                                    max_thread_count=max_thread_count,
                                    auth=auth,
                                    headers=headers)

    @staticmethod
    def _get_single(url: str,
                    on_receive: Callable[[int, bytes, int], None],
                    on_progress: Union[None, Callable[[int, int], None]] = None,
                    auth: HTTPBasicAuth = HTTPBasicAuth(None, None),
                    headers: Union[None, Dict[str, str]] = None) -> bool:
        try:
            response = requests.get(url,
                                    verify=False,
                                    auth=auth,
                                    stream=True,
                                    headers=headers)

            # update file_size from response
            file_size = -1
            if 'Content-Length' in response.headers:
                content_length = response.headers['Content-Length']
                if content_length.isnumeric():
                    file_size = int(content_length)

            # update progress for start
            offset: int = 0
            if on_progress:
                on_progress(offset, file_size)

            # do download
            for chunk in response.iter_content(Http.DOWNLOAD_BUFFER_SIZE):
                on_receive(offset, chunk, file_size)
                offset += len(chunk)

                if on_progress:
                    on_progress(offset, file_size)

            # eos
            on_receive(file_size, b'', file_size)

            return True
        except Exception as e:
            print(f'error downloading: {e}')
            import traceback
            traceback.print_exception(e)
            return False

    @staticmethod
    def _get_mt(url: str,
                file_size: int,
                on_receive: Callable[[int, bytes, int], None],
                on_progress: Union[None, Callable[[int, int], None]] = None,
                min_segment_size: int = 1 * 1024 * 1024,  # min segment size, 1M as default
                max_segment_size: int = 10 * 1024 * 1024,  # max segment size, 10M as default
                max_thread_count: int = 5,
                auth: HTTPBasicAuth = HTTPBasicAuth(None, None),
                headers: Union[None, Dict[str, str]] = None) -> bool:
        # get segment ranges
        drm = _DownloadRangeManager(0, file_size, min_segment_size, max_segment_size)

        # prepare context
        context = _MTContext(url,
                             auth,
                             file_size,
                             headers,
                             max_thread_count,
                             on_receive,
                             drm)

        # start segment downloading
        with ThreadPoolExecutor(max_workers=max_thread_count) as pool:
            futures = list()

            # add tasks
            futures.append(pool.submit(Http._on_add_tasks_mt, context, pool, futures))

            # update progress
            with context.lock:
                if on_progress:
                    on_progress(context.downloaded_size, file_size)

                while context.downloaded_size < file_size:
                    context.cond.wait()
                    if on_progress:
                        on_progress(context.downloaded_size, file_size)

            # eos
            on_receive(file_size, b'', file_size)

            as_completed(futures)

        return True

    @staticmethod
    def _on_add_tasks_mt(context: _MTContext, pool: ThreadPoolExecutor, futures: List):
        with context.lock:
            no_more_range: bool = False
            while context.downloading_tasks < context.max_thread_count:
                r = context.drm.get()
                if r.empty():
                    no_more_range = True
                    break

                futures.append(pool.submit(Http._on_get_mt, r, context))
                context.downloading_tasks += 1

            if not no_more_range:
                futures.append(pool.submit(Http._on_add_tasks_mt, context, pool, futures))

    @staticmethod
    def _on_get_mt(download_range: _Range, context: _MTContext):
        r = download_range.duplicate()

        try:
            # create header
            headers = {k: v for k, v in context.headers.items()}
            headers['Range'] = f'bytes={download_range.start}-{download_range.end}'  # overwrite range

            # send request
            response = requests.get(context.url, verify=False, auth=context.auth, stream=True, headers=headers)
            # content_range = response.headers['Content-Range']

            # do download
            for chunk in response.iter_content(min(Http.DOWNLOAD_BUFFER_SIZE, len(r))):
                with context.lock:
                    # commit
                    received_range = _Range.create_offset_length(r.start, len(chunk))
                    accepted_range = context.drm.commit(received_range)

                    # stop download if accepted range is empty
                    if accepted_range.empty():
                        if Http.DEBUG:
                            print(f'>>> completed(cut2): [{download_range.start} - {r.start - 1}]')
                        break  # drop data and stop this thread due to range already downloaded by other thread
                    accepted_chunk = chunk[0: len(accepted_range)]

                    # callback on_receive
                    context.on_receive(accepted_range.start, accepted_chunk, context.file_size)

                    # update downloaded_size and notify update progress
                    context.downloaded_size += len(accepted_chunk)
                    context.cond.notify_all()

                    # update next range
                    r = _Range(accepted_range.end + 1, r.end)
                    if r.empty():
                        if Http.DEBUG:
                            print(f'>>> completed(done): [{download_range.start} - {download_range.end}]')
                        break  # next range is empty. range fully downloaded
                    if len(accepted_range) < len(received_range):
                        if Http.DEBUG:
                            print(f'>>> completed(cut ): [{download_range.start} - {r.start - 1}]')
                        break  # accepted range is cut, it should reach the end
            response.close()

        except Exception as e:
            err_range = context.drm.abort(r)
            if Http.DEBUG:
                print(f'>>> completed(err ): [{download_range.start} - {err_range.start - 1}] {e}')

        # downloading tasks --
        with context.lock:
            context.downloading_tasks -= 1

    @staticmethod
    def get_headers(headers: Union[None, Dict[str, str]] = None) -> Dict[str, str]:
        return Http.HEADERS if headers is None else headers
